from pathlib import Path
from torch.utils.data import DataLoader
from dataset_utils import SpeechDataset,subsample2,subsample4
import torchaudio

###################################### Parameters of Speech processing ##################################
SAMPLE_RATE = 48000
N_FFT = 1022
HOP_LENGTH = 256
noise_class = "example_1"
if noise_class == "white":
    TRAIN_INPUT_DIR = Path('/home/abc/n2n/Datasets/WhiteNoise_Train_Input')
    TRAIN_TARGET_DIR = Path('/home/abc/n2n/Datasets/WhiteNoise_Train_Output')

    TEST_NOISY_DIR = Path('/home/abc/n2n/Datasets/WhiteNoise_Test_Input')
    TEST_CLEAN_DIR = Path('/home/abc/n2n/Datasets/clean_testset_wav')

# Load urbansound8K noise
else:
    TRAIN_INPUT_DIR = Path('D:\python_lab\denoiser\Only-Noisy-Training\Audio_example/' + str(noise_class))
    TRAIN_TARGET_DIR = Path('D:\python_lab\denoiser\Only-Noisy-Training\Audio_example/' + str(noise_class))

#列出所有wav音频文件的地址
train_input_files = sorted(list(TRAIN_INPUT_DIR.rglob('*.wav')))
train_target_files = sorted(list(TRAIN_TARGET_DIR.rglob('*.wav')))
#注意这里一定要加str
waveform, sample_rate = torchaudio.load(str(train_input_files[3]))
waveform = waveform.numpy()
print(waveform.shape)
print(sample_rate)
#加载数据集dataset
# #返回值分别是x_noisy_stft, g1_stft, g1_wav, g2_wav, x_clean_stft
# train_dataset = SpeechDataset(train_input_files, train_target_files, N_FFT, HOP_LENGTH)
# #将dataset包装为dataloader
# train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)